package com.yeziji.security.common.logger;

import javax.annotation.Nonnull;
import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.Objects;

/**
 * logger 捕获响应
 *
 * @author hwy
 * @since 2023/11/13 0:23
 **/
public class LoggerResponseWrapper extends HttpServletResponseWrapper {
    /**
     * 当前响应
     */
    private final HttpServletResponse response;

    /**
     * 响应体流
     */
    private final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();

    public LoggerResponseWrapper(HttpServletResponse response) {
        super(response);
        this.response = response;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return new PrintWriter(new OutputStreamWriter(this.byteArrayOutputStream, this.response.getCharacterEncoding()));
    }

    /**
     * 获取响应对象字符串
     *
     * @return {@link String} 响应对象字符串
     */
    public String getBody() {
        try {
            PrintWriter writer = this.getWriter();
            if (Objects.nonNull(writer)) {
                writer.flush();
                writer.close();
            }
            return this.byteArrayOutputStream.toString();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 可重复流
     */
    @Override
    public ServletOutputStream getOutputStream() {
        return new ServletOutputStream() {
            @Override
            public void write(int b) {
                byteArrayOutputStream.write(b);
            }

            @Override
            public void write(@Nonnull byte[] b) throws IOException {
                super.write(b);
            }

            @Override
            public void write(@Nonnull byte[] b, int off, int len) throws IOException {
                super.write(b, off, len);
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {

            }

            /**
             * 重写 flush, 让流内容重新填充回 response 当中
             */
            @Override
            public void flush() throws IOException {
                if (!response.isCommitted()) {
                    byte[] body = byteArrayOutputStream.toByteArray();
                    ServletOutputStream outputStream = response.getOutputStream();
                    outputStream.write(body);
                    outputStream.flush();
                }
            }
        };
    }
}
